package com.chao.java.ai.langchain4j.store;

import com.chao.java.ai.langchain4j.bean.ChatMessages;
import com.github.xiaoymin.knife4j.core.util.StrUtil;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;

/**
 * @author eclipse
 * @version 2025/6/28 22:05
 */
@Component
@RequiredArgsConstructor
public class MongoChatMemoryStore implements ChatMemoryStore {

    private final MongoTemplate mongoTemplate;


    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);
        ChatMessages one = mongoTemplate.findOne(query, ChatMessages.class);
        if (Objects.isNull(one)) {
            return new LinkedList<>();
        }
        String content = one.getContent();
        if (StrUtil.isBlank(content)) {
            return new LinkedList<>();
        }
        return ChatMessageDeserializer.messagesFromJson(content);
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> list) {
        Query query = new Query(Criteria.where("memoryId").is(memoryId));
        Update update = new Update().set("content", ChatMessageSerializer.messagesToJson(list));
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        Query query = new Query(Criteria.where("memoryId").is(memoryId));
        mongoTemplate.remove(query, ChatMessages.class);
    }
}
